# coding=utf-8
# @author:      ChengJing
# @name:        IRNN.py
# @datetime:    2022/1/29 21:17
# @software:    PyCharm
# @description:

import torch
import torch.nn as nn
from model.improve_location.ImproveFeatures import LAQ

torch.manual_seed(66)
torch.cuda.manual_seed(66)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


class IRNN(nn.Module):
    """
    使用GRU模型构建RNN网络，进行爆管预警和位置的识别
    """
    def __init__(self, g, in_features, gru_hidden, linear_hidden, laq_hidden, out_features):
        """
        Args:
            g: 流量监测点和压力监测点之间的连接关系，shape：N*M
            in_features: 输入数据的特征
            gru_hidden: gru层的隐含层节点数
            linear_hidden: linear层的隐含层节点数
            laq_hidden: 隐含层节点数
            out_features: 输出数据的特征
        """
        super(IRNN, self).__init__()
        self.laq1 = LAQ(g, in_features, laq_hidden)
        self.gru = nn.GRU(
            input_size=in_features,
            hidden_size=gru_hidden,
            num_layers=3,
            # dropout=0.5,
            batch_first=True)
        # self.laq2 = LAQ(g, gru_hidden, laq_hidden)
        self.classification = nn.Sequential(
            nn.Linear(gru_hidden, linear_hidden),
            nn.ReLU(),
            # nn.Linear(linear_hidden, 2 * linear_hidden),
            # nn.Linear(2 * linear_hidden, linear_hidden),
            nn.Linear(linear_hidden, out_features),
            # nn.Linear(gru_hidden, out_features),
            nn.Softmax(dim=1)
        )

    def forward(self, data):
        """
        前向传递函数
        """
        x, q = data
        x = self.laq1(x, q)
        gru_x, _ = self.gru(x)
        # gru_x = self.laq2(gru_x, q)
        lx = gru_x[:, -1, :]
        px = self.classification(lx)
        return px


if __name__ == '__main__':
    x = torch.rand((10, 20, 6))
    q = torch.rand((10, 3))
    g = torch.rand((3, 6))
    model = IRNN(g, 6, 36, 72, 36, 20)
    y = model((x, q))
    print(y.shape)
    z = y.sum()
    z.requires_grad_(True)
    z.backward()
